# brute_force_demo.py
import time
import secrets

# 导入你的加密算法实现和辅助工具
from publickey.elgamal_core import ElGamal
from publickey.ecc_core import get_curve_by_name, generate_keypair as ecc_gen_keys, ecc_encrypt_simplified, ecc_decrypt_simplified
# 修改：从你自己的math_utils导入基础函数
from publickey.math_utils import generate_prime_number, extended_gcd, mod_inverse, power

# ==============================================================================
# 新增: 一个不带OAEP填充的“教科书式RSA”类，专门用于本次漏洞演示
# ==============================================================================
class TextbookRSA:
    """
    一个不使用任何填充的“教科书式”RSA实现，用于演示漏洞。
    警告：绝对不能在实际应用中使用此类实现！
    """
    def __init__(self, bits=1024):
        p = generate_prime_number(bits // 2)
        q = generate_prime_number(bits // 2)
        while p == q:
            q = generate_prime_number(bits // 2)
        
        self.n = p * q
        phi_n = (p - 1) * (q - 1)
        self.e = 65537
        while extended_gcd(self.e, phi_n)[0] != 1:
            self.e = 2 * secrets.randbelow((phi_n - 3) // 2) + 3 # 备用e
        
        self.d = mod_inverse(self.e, phi_n)
        self.public_key = (self.n, self.e)
        self.private_key = (self.n, self.d)

    def encrypt(self, message_int: int) -> int:
        """加密一个整数"""
        if not (0 <= message_int < self.n):
            raise ValueError("消息必须在 [0, n-1) 范围内")
        return power(message_int, self.e, self.n)

    def decrypt(self, ciphertext_int: int) -> int:
        """解密一个整数"""
        return power(ciphertext_int, self.d, self.n)

# ==============================================================================
# 1. 针对RSA的穷举攻击模拟 (修改后)
# ==============================================================================
def simulate_rsa_brute_force():
    """
    演示: 对一个密钥空间极小的RSA进行穷举攻击。
    """
    print("\n" + "="*20, "模拟: RSA私钥穷举攻击", "="*20)
    
    toy_rsa = TextbookRSA(bits=24) 
    n, e = toy_rsa.public_key
    real_d = toy_rsa.private_key[1]
    
    print(f"创建了一个玩具RSA (n={n}, e={e})。")
    print(f"真实的私钥 d (基于 phi(n)) 是: {real_d} (攻击者不知道此值)")

    known_plaintext = b"OK"
    known_plaintext_int = int.from_bytes(known_plaintext, 'big')
    ciphertext = toy_rsa.encrypt(known_plaintext_int)
    print(f"用公钥加密明文 '{known_plaintext.decode()}' 得到密文: {ciphertext}")

    print("攻击开始... 正在穷举尝试所有可能的d值...")
    start_time = time.perf_counter()
    
    for guessed_d in range(1, n):
        decrypted_attempt = pow(ciphertext, guessed_d, n)
        if decrypted_attempt == known_plaintext_int:
            end_time = time.perf_counter()
            print("\n" + "="*20, "攻击成功！", "="*20)
            print(f"++ 找到一个有效的私钥 d: {guessed_d}")
            print(f"++ 耗时: {end_time - start_time:.4f} 秒")
            # 修改：移除 assert guessed_d == real_d
            # 因为穷举找到的密钥(基于lambda(n))可能与计算出的密钥(基于phi(n))不同，但功能等价。
            # 能够成功解密，本身就证明了攻击成功。
            if guessed_d != real_d:
                print(f"注意: 找到的私钥({guessed_d})与原始私钥({real_d})不同，但同样有效！")
            return
    
    print("!! 攻击失败，未能找到私钥。")

# ... (ElGamal 和 ECC 的模拟函数保持不变) ...
def simulate_elgamal_brute_force():
    """演示: 对一个密钥空间极小的ElGamal进行穷举攻击。"""
    print("\n" + "="*20, "模拟: ElGamal私钥穷举攻击", "="*20)
    
    toy_elgamal = ElGamal(key_size_bits=24)
    p, _, g, y = toy_elgamal.public_key
    real_x = toy_elgamal.private_key
    
    print(f"创建了一个玩具ElGamal (p={p})。")
    print(f"真实的私钥 x 是: {real_x} (攻击者不知道此值)")
    
    known_plaintext_int = 12345
    c1, c2 = ElGamal.encrypt(toy_elgamal.public_key, known_plaintext_int)
    print(f"用公钥加密明文 '{known_plaintext_int}' 得到密文: ({c1}, {c2})")
    
    print("攻击开始... 正在穷举尝试所有可能的x值...")
    start_time = time.perf_counter()
    
    # 在ElGamal中，私钥是唯一的
    for guessed_x in range(1, p):
        s = pow(c1, guessed_x, p)
        s_inv = pow(s, -1, p)
        decrypted_attempt = (c2 * s_inv) % p
        if decrypted_attempt == known_plaintext_int:
            end_time = time.perf_counter()
            print("\n" + "="*20, "攻击成功！", "="*20)
            print(f"++ 找到私钥 x: {guessed_x}")
            print(f"++ 耗时: {end_time - start_time:.4f} 秒")
            assert guessed_x == real_x
            return
            
    print("!! 攻击失败，未能找到私钥。")

def simulate_ecc_brute_force():
    """演示: 对一个密钥空间极小的ECC进行穷举攻击。"""
    print("\n" + "="*20, "模拟: ECC私钥穷举攻击", "="*20)
    
    curve = get_curve_by_name("secp256k1")
    search_space_limit = 2**22
    
    real_d_small = secrets.randbelow(search_space_limit - 1) + 1
    pub_key_from_small_d = curve.scalar_multiply(real_d_small, curve.G)
    
    print(f"创建了一个ECC密钥对，其私钥被限制在一个很小的空间内。")
    print(f"真实的私钥 d 是: {real_d_small} (攻击者不知道此值)")
    
    known_plaintext = b"test"
    R, C = ecc_encrypt_simplified(pub_key_from_small_d, curve, known_plaintext)
    print(f"用公钥加密明文 '{known_plaintext.decode()}' 得到密文。")
    
    print(f"攻击开始... 正在穷举尝试 [1, {search_space_limit-1}] 范围内的d值...")
    start_time = time.perf_counter()
    
    # 在ECC中，私钥也是唯一的
    for guessed_d in range(1, search_space_limit):
        guessed_pub_key = curve.scalar_multiply(guessed_d, curve.G)
        if guessed_pub_key == pub_key_from_small_d:
            end_time = time.perf_counter()
            decrypted = ecc_decrypt_simplified(guessed_d, curve, R, C)
            assert decrypted == known_plaintext
            print("\n" + "="*20, "攻击成功！", "="*20)
            print(f"++ 找到私钥 d: {guessed_d}")
            print(f"++ 耗时: {end_time - start_time:.4f} 秒")
            return

    print("!! 攻击失败，未能找到私钥。")


if __name__ == '__main__':
    print("#"*20, " 通用攻击：密钥空间穷举模拟 ", "#"*20)
    print("本系列实验旨在证明，密码系统的安全性根本上依赖于其巨大的密钥空间。")
    print("通过在“玩具级”的小密钥空间上进行穷举，我们可以直观地看到这种防御是如何被瓦解的。")
    
    simulate_rsa_brute_force()
    simulate_elgamal_brute_force()
    simulate_ecc_brute_force()